
namespace Microsoft.VsSDK.IntegrationTestLibrary
{
	using System;
	using System.Collections.Generic;
	using System.Text;
	using System.Runtime.InteropServices;
	using System.Threading;
	using Microsoft.VisualStudio.Shell.Interop;
	using Microsoft.VisualStudio.Shell;

	/// <summary>
	/// This class is responsible to close dialog boxes that pop up during different VS Calls
	/// </summary>
	internal class DialogBoxPurger : IDisposable
	{
		/// <summary>
		/// The default number of milliseconds to wait for the threads to signal to terminate.
		/// </summary>
		private const int DefaultMillisecondsToWait = 3500;

		/// <summary>
		/// Object used for synchronization between thread calls.
		/// </summary>
		internal static volatile object Mutex = new object();

		/// <summary>
		/// The IVsUIShell. This cannot be queried on the working thread from the service provider. Must be done in the main thread.!!
		/// </summary>
		private IVsUIShell uiShell;

		/// <summary>
		/// The button to "press" on the dialog.
		/// </summary>
		private int buttonAction;

		/// <summary>
		/// Thread signales to the calling thread that it is done.
		/// </summary>
		private bool exitThread = false;

		/// <summary>
		/// Calling thread signales to this thread to die.
		/// </summary>
		private AutoResetEvent threadDone = new AutoResetEvent(false);

		/// <summary>
		/// The queued thread started.
		/// </summary>
		private AutoResetEvent threadStarted = new AutoResetEvent(false);

		/// <summary>
		/// The result of the dialogbox closing for all the dialog boxes. That is if there are two of them and one fails this will be false.
		/// </summary>
		private bool dialogBoxCloseResult = false;

		/// <summary>
		/// The expected text to see on the dialog box. If set the thread will continue finding the dialog box with this text.
		/// </summary>
		private string expectedDialogBoxText = String.Empty;

		/// <summary>
		/// The number of the same  dialog boxes to wait for.
		/// This is for scenarios when two dialog boxes with the same text are popping up.
		/// </summary>
		private int numberOfDialogsToWaitFor = 1;

		/// <summary>
		/// Has the object been disposed.
		/// </summary>
		private bool isDisposed;

		/// <summary>
		/// Overloaded ctor.
		/// </summary>
		/// <param name="buttonAction">The botton to "press" on the dialog box.</param>
		/// <param name="numberOfDialogsToWaitFor">The number of dialog boxes with the same message to wait for. This is the situation when the same action pops up two of the same dialog boxes</param>
		/// <param name="expectedDialogMesssage">The expected dialog box message to check for.</param>
		internal DialogBoxPurger(int buttonAction, int numberOfDialogsToWaitFor, string expectedDialogMesssage)
		{
			this.buttonAction = buttonAction;
			this.numberOfDialogsToWaitFor = numberOfDialogsToWaitFor;
			this.expectedDialogBoxText = expectedDialogMesssage;
		}

		/// <summary>
		/// Overloaded ctor.
		/// </summary>
		/// <param name="buttonAction">The botton to "press" on the dialog box.</param>
		/// <param name="numberOfDialogsToWaitFor">The number of dialog boxes with the same message to wait for. This is the situation when the same action pops up two of the same dialog boxes</param>
		internal DialogBoxPurger(int buttonAction, int numberOfDialogsToWaitFor)
		{
			this.buttonAction = buttonAction;
			this.numberOfDialogsToWaitFor = numberOfDialogsToWaitFor;
		}

		/// <summary>
		/// Overloaded ctor.
		/// </summary>
		/// <param name="buttonAction">The botton to "press" on the dialog box.</param>
		/// <param name="expectedDialogMesssage">The expected dialog box message to check for.</param>
		internal DialogBoxPurger(int buttonAction, string expectedDialogMesssage)
		{
			this.buttonAction = buttonAction;
			this.expectedDialogBoxText = expectedDialogMesssage;
		}

		/// <summary>
		/// Overloaded ctor.
		/// </summary>
		/// <param name="buttonAction">The botton to "press" on the dialog box.</param>
		internal DialogBoxPurger(int buttonAction)
		{
			this.buttonAction = buttonAction;
		}

		/// <summary>
		#region IDisposable Members

		void IDisposable.Dispose()
		{
			if (this.isDisposed)
			{
				return;
			}

			this.WaitForDialogThreadToTerminate();

			this.isDisposed = true;
		}

		/// <summary>
		/// Spawns a thread that will start listening to dialog boxes.
		/// </summary>
		internal void Start()
		{
			// We ask for the uishell here since we cannot do that on the therad that we will spawn.
			IVsUIShell uiShell = Package.GetGlobalService(typeof(SVsUIShell)) as IVsUIShell;

			if (uiShell == null)
			{
				throw new InvalidOperationException("Could not get the uiShell from the serviceProvider");
			}

			this.uiShell = uiShell;

			System.Threading.Thread thread = new System.Threading.Thread(new ThreadStart(this.HandleDialogBoxes));
			thread.Start();

			// We should never deadlock here, hence do not use the lock. Wait to be sure that the thread started.
			this.threadStarted.WaitOne(3500, false);
		}

		/// <summary>
		/// Waits for the dialog box close thread to terminate. If the thread does not signal back within millisecondsToWait that it is shutting down,
		/// then it will tell to the thread to do it.
		/// </summary>
		internal bool WaitForDialogThreadToTerminate()
		{
			return this.WaitForDialogThreadToTerminate(DefaultMillisecondsToWait);
		}

		/// <summary>
		/// Waits for the dialog box close thread to terminate. If the thread does not signal back within millisecondsToWait that it is shutting down,
		/// then it will tell to the thread to do it.
		/// </summary>
		/// <param name="millisecondsToWait">The number milliseconds to wait for until the dialog purger thread is signaled to terminate. This is just for safe precaution that we do not hang. </param>
		/// <returns>The result of the dialog boxes closing</returns>
		internal bool WaitForDialogThreadToTerminate(int numberOfMillisecondsToWait)
		{
			bool signaled = false;

			// We give millisecondsToWait sec to bring up and close the dialog box.
			signaled = this.threadDone.WaitOne(numberOfMillisecondsToWait, false);

			// Kill the thread since a timeout occured.
			if (!signaled)
			{
				lock (Mutex)
				{
					// Set the exit thread to true. Next time the thread will kill itselfes if it sees 
					this.exitThread = true;
				}

				// Wait for the thread to finish. We should never deadlock here.
				this.threadDone.WaitOne();
			}

			return this.dialogBoxCloseResult;
		}

		/// <summary>
		/// This is the thread method. 
		/// </summary>
		private void HandleDialogBoxes()
		{
			// No synchronization numberOfDialogsToWaitFor since it is readonly
			IntPtr[] hwnds = new IntPtr[this.numberOfDialogsToWaitFor];
			bool[] dialogBoxCloseResults = new bool[this.numberOfDialogsToWaitFor];

			try
			{
				// Signal that we started
				lock (Mutex)
				{
					this.threadStarted.Set();
				}

				// The loop will be exited either if a message is send by the caller thread or if we found the dialog. If a message box text is specified the loop will not exit until the dialog is found.
				bool stayInLoop = true;
				int dialogBoxesToWaitFor = 1;

				while (stayInLoop)
				{
					int hwndIndex = dialogBoxesToWaitFor - 1;

					// We need to lock since the caller might set context to null.
					lock (Mutex)
					{
						if (this.exitThread)
						{
							break;
						}

						// We protect the shell too from reentrency.
						this.uiShell.GetDialogOwnerHwnd(out hwnds[hwndIndex]);

					}

					if (hwnds[hwndIndex] != IntPtr.Zero)
					{
						StringBuilder windowClassName = new StringBuilder(256);
						NativeMethods.GetClassName(hwnds[hwndIndex], windowClassName, windowClassName.Capacity);

						// The #32770 is the class name of a messagebox dialog.
						if (windowClassName.ToString().Contains("#32770"))
						{
							IntPtr unmanagedMemoryLocation = IntPtr.Zero;
							string dialogBoxText = String.Empty;
							try
							{
								unmanagedMemoryLocation = Marshal.AllocHGlobal(10 * 1024);
								NativeMethods.EnumChildWindows(hwnds[hwndIndex], new NativeMethods.CallBack(FindMessageBoxString), unmanagedMemoryLocation);
								dialogBoxText = Marshal.PtrToStringUni(unmanagedMemoryLocation);
							}
							finally
							{
								if (unmanagedMemoryLocation != IntPtr.Zero)
								{
									Marshal.FreeHGlobal(unmanagedMemoryLocation);
								}
							}

							lock (Mutex)
							{

								// Since this is running on the main thread be sure that we close the dialog.
								bool dialogCloseResult = false;
								if (this.buttonAction != 0)
								{
									dialogCloseResult = NativeMethods.EndDialog(hwnds[hwndIndex], this.buttonAction);
								}

								// Check if we have found the right dialog box.
								if (String.IsNullOrEmpty(this.expectedDialogBoxText) || (!String.IsNullOrEmpty(dialogBoxText) && String.Compare(this.expectedDialogBoxText, dialogBoxText.Trim(), StringComparison.OrdinalIgnoreCase) == 0))
								{
									dialogBoxCloseResults[hwndIndex] = dialogCloseResult;
									if (dialogBoxesToWaitFor++ >= this.numberOfDialogsToWaitFor)
									{
										stayInLoop = false;
									}
								}
							}
						}
					}
				}
			}
			finally
			{
				//Let the main thread run a possible close command.
				System.Threading.Thread.Sleep(2000);

				foreach (IntPtr hwnd in hwnds)
				{
					// At this point the dialog should be closed, if not attempt to close it.
					if (hwnd != IntPtr.Zero)
					{
						NativeMethods.SendMessage(hwnd, NativeMethods.WM_CLOSE, 0, new IntPtr(0));
					}
				}

				lock (Mutex)
				{
					// Be optimistic.
					this.dialogBoxCloseResult = true;

					for (int i = 0; i < dialogBoxCloseResults.Length; i++)
					{
						if (!dialogBoxCloseResults[i])
						{
							this.dialogBoxCloseResult = false;
							break;
						}
					}

					this.threadDone.Set();
				}
			}
		}

		/// <summary>
		/// Finds a messagebox string on a messagebox.
		/// </summary>
		/// <param name="hwnd">The windows handle of the dialog</param>
		/// <param name="unmanagedMemoryLocation">A pointer to the memorylocation the string will be written to</param>
		/// <returns>True if found.</returns>
		private static bool FindMessageBoxString(IntPtr hwnd, IntPtr unmanagedMemoryLocation)
		{
			StringBuilder sb = new StringBuilder(512);
			NativeMethods.GetClassName(hwnd, sb, sb.Capacity);

			if (sb.ToString().ToLower().Contains("static"))
			{
				StringBuilder windowText = new StringBuilder(2048);
				NativeMethods.GetWindowText(hwnd, windowText, windowText.Capacity);

				if (windowText.Length > 0)
				{
					IntPtr stringAsPtr = IntPtr.Zero;
					try
					{
						stringAsPtr = Marshal.StringToHGlobalAnsi(windowText.ToString());
						char[] stringAsArray = windowText.ToString().ToCharArray();

						// Since unicode characters are copied check if we are out of the allocated length.
						// If not add the end terminating zero.
						if ((2 * stringAsArray.Length) + 1 < 2048)
						{
							Marshal.Copy(stringAsArray, 0, unmanagedMemoryLocation, stringAsArray.Length);
							Marshal.WriteInt32(unmanagedMemoryLocation, 2 * stringAsArray.Length, 0);
						}
					}
					finally
					{
						if (stringAsPtr != IntPtr.Zero)
						{
							Marshal.FreeHGlobal(stringAsPtr);
						}
					}
					return false;
				}
			}

			return true;
		}

		#endregion
	}
}
